# coding:utf-8

import matplotlib.pyplot as plt
import io
import os
import sys
import numpy as np
import fitz  # pdf-img
import base64
from PIL import Image
import cv2 as cv
from skimage import io as skio, morphology
import glob
import pytesseract


# files
def addFile(inFilePath):
    flag = 0
    if os.path.exists(inFilePath):
        flag = 1
        print(f'{inFilePath} has already existed.\n backup in {inFilePath}.bak')
        if os.path.exists(f"{inFilePath}.bak"):
            os.remove(f"{inFilePath}.bak")
        with open(inFilePath + '.bak', "w")as nf, open(inFilePath, 'r')as rf:
            nf.write(rf.read())
        os.remove(inFilePath)
    with open(inFilePath, "w") as f:
        pass
    return flag


def moveFile(inFilePath, outPath, backUp=True):
    flag = copyFile(inFilePath, outPath, backUp)
    if flag == 0:
        os.remove(inFilePath)


def copyFile(inFilePath, outPath, backUp=True):
    baseName = os.path.basename(inFilePath)
    OutIsDir = os.path.basename(outPath) == ''
    if OutIsDir:
        outFilePath = os.path.join(outPath, baseName)
        # 备份
        if os.path.exists(outFilePath) and backUp:
            print(f'{outFilePath} has already existed.\n backup in {outFilePath}.bak')
            with open(outFilePath + '.bak', "w")as nf, open(outFilePath, 'r')as rf:
                nf.write(rf.read())
        os.remove(outFilePath)
        with open(outFilePath, 'w')as nf2, open(inFilePath, 'r')as rf2:
            nf2.write(rf2.read())
        print(f"moved {inFilePath} , to {outFilePath}")
        return 0
    else:
        print('outPath not dir, failed to move file')
        return 1


def deletFile(inFilePath):
    if os.path.exists(inFilePath):
        os.remove(inFilePath)
    return


def readPdf(inPath):
    """
    PDF转换成dnarray图像
    返回一个pdf页数维度的 np.ndarray 数组
    本步骤占用较多内存，单个PDF的页数不要过多
    This step takes up a lot of RAM, one PDF file should not have too many pages.
    """
    with fitz.open(inPath) as pdf:
        # 初始化数组
        pageNum = pdf.pageCount
        pageShape = np.array(Image.open(io.BytesIO(
            pdf[0].getPixmap(matrix=fitz.Matrix(25 / 9, 25 / 9).preRotate(0), alpha=False).getPNGData())).convert(
            "L")).shape
        imgList = np.empty(pageNum * pageShape[0] * pageShape[1]).reshape(pageNum, pageShape[0], pageShape[1])
        # 逐页读取PDF

        for i in range(pageNum):
            page = pdf[i]
            print('read in ', end='')
            print(page)
            # 扩大倍数变得清晰，不旋转。选择25/9是为了和WPS工具保持一致
            trans = fitz.Matrix(25 / 9, 25 / 9).preRotate(0)
            p = page.getPixmap(matrix=trans, alpha=False).getPNGData()
            print(type(p))
            # p是bytes(base64),以下用PIL和io
            img = np.array(Image.open(io.BytesIO(p)).convert("L"))
            # plt.imshow(img)
            # plt.show()
            # saveImg(f'./img{i}.png',img)
            imgList[i] = img
    return imgList


def readImg(inImgPath):
    img = skio.imread(inImgPath, as_gray=True)
    return img


def saveImg(outImgPath, img):
    baseName = os.path.basename(outImgPath)
    if baseName != '':
        skio.imsave(os.path.join(outImgPath), img)
        return 0
    else:
        print(f"{outImgPath} not a img name")
        return 1


def binary(img, bth: float):
    """
    bth 阈值，将灰度图转化为二值图
    """
    img[img < bth] = 0
    img[img >= bth] = 1
    return img


def read_binary(inImgPath, bth=0.6):
    """
    为了方便使用，这个函数可以读取不同格式文档，推荐使用
    """
    filetype = os.path.basename(inImgPath).split('.')[-1]
    if filetype.lower() == "pdf":
        imglist = readPdf(inImgPath)
        for i in range(imglist.shape[0]):
            # saveImg(f'.binary_before{i}.png', imglist[i])
            temp_img = binary(imglist[i], bth * 255)
            imglist[i] = temp_img
            # saveImg(f'.binary{i}.png',imglist[i])
        return imglist
    if filetype.lower() in ['png', 'jpg', 'jpeg']:
        return binary(readImg(inImgPath), bth)


def correct(img):
    # todo
    """
    这段代码成立的假设前提是：
    要处理的非标准表格的第一列第一个点和最后一列第一个点在同一行
    表格第一列和最后一列是垂直的，最上方和最下方的交点水平距离的不超过阈值（默认20像素）
    """
    points = getCrossPoint(img)
    # 分别获取第一列和最后一列，筛选出水平距离不超过阈值的点位，获取行序号最小的一个（最上方的点）
    MAX_TH = 20
    temp = points[:, 1] - points[:, 1].min()
    fTop = np.sort(points[np.argwhere(temp < MAX_TH)], axis=0)[0]
    temp2 = points[:, 1].max() - points[:, 1]
    lTop = np.sort(points[np.argwhere(temp2 < MAX_TH)], axis=0)[0]
    # 两个点坐标插值
    d = (fTop - lTop)[0]
    # 判断倾斜方向，如果第一个点位低，右方的图像坐标变为原坐标+偏移量，向下移动
    flag = 1 if d[0] > 0 else -1
    # 根据局部判断整体的形变程度
    H, W = img.shape
    dy, dx = d * W // d[1]
    # 把原图进行上下平移
    steps = abs(dy)
    partW = W // steps
    for i in range(steps):
        stCol = i * partW
        endCol = stCol + partW if i != steps - 1 else W
        partimg = img[i:H - i, stCol:endCol]
        img[i + flag * i:H - i + flag * i, stCol:endCol] = partimg
    saveImg("./corrected.png", img)
    return img


def opening(img, selemShape=(3, 3)):
    # Opening (morphology),dilation + erosion
    # 先膨胀后腐蚀，开运算去噪，这里是祛除文字，
    # selem为结构元素
    w, h = selemShape
    selem = morphology.rectangle(w, h)
    img = morphology.dilation(img, selem)
    img = morphology.erosion(img, selem)
    return img


def closing(img, selemShape=(3, 3)):
    # Closing (morphology),dilation + erosion
    # 先腐蚀后膨胀，闭运算链接，这里是链接断线，
    # selem为结构元素
    w, h = selemShape
    selem = morphology.rectangle(w, h)
    img = morphology.erosion(img, selem)
    img = morphology.dilation(img, selem)
    return img


def getCrossPoint(img):
    # 返回交叉点的坐标信息
    rowSelemShape = (1, img.shape[0] // 60)
    colSelemShape = (img.shape[1] // 60, 1)
    horizontalLine = opening(img, rowSelemShape)
    verticalLine = opening(img, colSelemShape)
    verticalLine = closing(verticalLine, (1, verticalLine.shape[1] // 20))
    # 获取交点
    crossPoints = (1 - horizontalLine) * (1 - verticalLine)
    # saveImg('./crosspoingboforecheck.png', crossPoints)
    crossPosition = np.argwhere(crossPoints)
    # 删除点附近曼哈顿距离小于9的点的位置（保留的是最后一个点）
    for i in range(crossPosition.shape[0]):
        d = np.sum(np.abs(crossPosition - crossPosition[i, :]), axis=1)
        nearBy = np.argwhere(d < 9)
        itSelf = np.argwhere(d == 0)
        for point in nearBy:
            r, c = crossPosition[point[0]]
            r0, c0 = crossPosition[itSelf[0][0]]
            crossPoints[r, c] = 0
            crossPoints[r0, c0] = 1
    points = np.argwhere(crossPoints) # 交点的位置（乱序）
    return points

def getCrossPoints_sepByRow(img):
    # 为了方便使用，这里讲返回的序列按照行进行了重新排列，返回一个列表
    points=getCrossPoint(img)
    sortedCrossPoints = np.sort(points, axis=0)
    d = np.diff(sortedCrossPoints)
    # 通过插值法获取第x行的交点坐标，插值超过12被认定为新行
    rowEnd = np.argwhere(d > 12)
    rowList = []
    for i in range(len(rowEnd) - 1):
        temp = sortedCrossPoints[rowEnd[i]:rowEnd[i + 1]]
        rowList.append(np.sort(temp, axis=1))
    return rowList

def cut(img, basePoint=(0, 0), shape=(256, 256)):
    childImg = img[basePoint[0]:basePoint[0] + shape[0], basePoint[1]:basePoint[1] + shape[1]]
    return childImg


def setModel(modelPath):
    # todo
    print(f'using model {modelPath}')
    pass


def ocr(img, language):
    # todo
    language=language.lower()
    langDic={
        "chinese":"chi_sim",
        "english":"eng"
    }
    if language=="":
        text=pytesseract.image_to_string(img)
    else:
        text=pytesseract.image_to_string(img,lang=langDic[language])
    return text



if __name__ == '__main__':
    img = read_binary("X:\Python\IGS\data_set\宋园路7月打印.pdf")[1]
    saveImg('./img.png', img)
    correct(img)
    pass
